# spark_api.py
# coding: utf-8
import json
import base64
import hashlib
import hmac
import ssl
from datetime import datetime
from time import mktime
from urllib.parse import urlencode, urlparse
import websocket
import _thread as thread
from wsgiref.handlers import format_date_time
from api_base import APIBase

class SparkAPI(APIBase):
    def __init__(self, config):
        super().__init__(config)
        self.generated_content = None
        self.parameter = None
        self.content_buffer = []  # 新增内容缓冲区

    def _generate_parameter(self, topic, role_content):
        return {
            "header": {"app_id": self.config['APP_ID']},
            "parameter": {
                "chat": {
                    "domain": "4.0Ultra",
                    "temperature": 0.5,
                    "max_tokens": 4096,
                    "top_k": 4,
                    "tools": [
                        {
                            "type": "web_search",
                            "web_search": {
                                "enable": True,
                                "search_mode": "deep",
                            }
                        }
                    ]
                }
            },
            "payload": {
                "message": {
                    "text": [
                        {
                            "role": "system",
                            "content": role_content
                        },
                        {
                            "role": "user",
                            "content": topic
                        }
                    ]
                }
            }
        }

    def _on_message(self, ws, message):
        data = json.loads(message)
        if data['header']['code'] != 0:
            print(f"API错误: {data['header']['message']}")
            ws.close()
        else:
            content_chunk = data['payload']['choices']['text'][0]['content']
            self.content_buffer.append(content_chunk)
            if data['payload']['choices']['status'] == 2:
                raw_content = "".join(self.content_buffer)
                formatted_content = '\n'.join([f'<p>{p.strip()}</p>' for p in raw_content.split('\n') if p.strip()])
                self.generated_content = formatted_content
                self.content_buffer = []
                ws.close()

    @staticmethod
    def _on_error(ws, error):
        print("### 错误:", error)

    @staticmethod
    def _on_close(ws, close_status_code, close_msg):
        print("### 连接关闭 ###")

    def _on_open(self, ws):
        def run(*args):
            data = json.dumps(self.parameter)
            ws.send(data)

        thread.start_new_thread(run, ())

    def generate_news(self, topic, role_content):
        self.generated_content = None
        self.parameter = self._generate_parameter(topic, role_content)  # 传递角色描述
        ws_param = WsParam(self.config['API_KEY'], self.config['API_SECRET'], self.config['API_URL'])
        ws_url = ws_param.create_url()
        ws = websocket.WebSocketApp(
            ws_url,
            on_open=self._on_open,
            on_message=self._on_message,
            on_error=self._on_error,
            on_close=self._on_close
        )
        ws.run_forever(sslopt={"cert_reqs": ssl.CERT_NONE})
        return self.generated_content


class WsParam:
    def __init__(self, api_key, api_secret, gpt_url):
        self.api_key = api_key
        self.api_secret = api_secret
        self.host = urlparse(gpt_url).netloc
        self.path = urlparse(gpt_url).path
        self.gpt_url = gpt_url

    def create_url(self):
        now = datetime.now()
        date = format_date_time(mktime(now.timetuple()))
        signature_origin = f"host: {self.host}\ndate: {date}\nGET {self.path} HTTP/1.1"
        signature_sha = hmac.new(
            self.api_secret.encode('utf-8'),
            signature_origin.encode('utf-8'),
            digestmod=hashlib.sha256
        ).digest()
        signature_sha_base64 = base64.b64encode(signature_sha).decode(encoding='utf-8')
        authorization_origin = f'api_key="{self.api_key}", algorithm="hmac-sha256", headers="host date request-line", signature="{signature_sha_base64}"'.encode('utf-8')
        authorization = base64.b64encode(authorization_origin).decode('utf-8')
        v = {"authorization": authorization, "date": date, "host": self.host}
        url = self.gpt_url + '?' + urlencode(v)
        return url
